package com.nmm.demo.controller;

import jakarta.annotation.Resource;
import java.util.List;
import java.util.stream.Collectors;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.prompt.PromptTemplate;
import org.springframework.ai.document.Document;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.util.ObjectUtils;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;

/**
 * 机器人对话控制器
 * @author niemingming
 * @Date 2025/6/12
 */
@RestController
@RequestMapping("/open/ai/v1")
public class ChatAiController {

  @Resource(name = "ollamaChatClient")
  private ChatClient ollamaChatClient;

  @Resource(name = "businessChatClient")
  private ChatClient businessChatClient;
  @Resource
  private VectorStore vectorStore;

  @GetMapping("/chat")
  public String chatMessage(@RequestParam(name = "question") String question) {
    boolean flag = question.startsWith("查询业务:");
    question = question.substring(5);
    if (flag) {
      return businessChatClient.prompt().user(question).call().content();
    }
    List<Document> docs = vectorStore.similaritySearch(SearchRequest.builder()
            .topK(3)
        .similarityThreshold(0.3) // 相识度阈值
        .query(question)
        .build());
    System.out.println(docs);


    return ollamaChatClient.prompt(getDocumentPromptTemplate(docs, question).create())
        .user(question).call().content();
  }

  /**
   * [获取提示词模板]
   * @param docs 查询到的文档信息
   * @param question 问题信息
   * @author niemingming 2025/6/12
   */
  private PromptTemplate getDocumentPromptTemplate(List<Document> docs, String question) {
    String template = """
        请基于以下上下文回答问题：
        {documents}
        
        用户问题：{question}
        """;
    PromptTemplate promptTemplate = new PromptTemplate(template);
    promptTemplate.add("documents", docs.stream().map(Document::getText).collect(Collectors.joining("\n")));
    promptTemplate.add("question", question);
    return promptTemplate;
  }

}
